
from wheezy.captcha.image import (
    background,
    captcha,
    # curve,
    noise,
    offset,
    rotate,
    smooth,
    text,
    warp,
)


class Captcha_Product:
    def __init__(self, width: int = 150,
                 height: int = 150,
                 *,
                 big_char: bool = False,
                 small_char: bool = False,
                 mul_num: int = 1,
                 ):
        """
        验证码生成器

        :param big_char: 是否添加大写字母
        :param small_char: 是否添加小写字母
        :param mul_num: 倍数
        """
        self.width = width
        self.height = height
        self.choice_list = [str(i) for i in range(10)] * mul_num
        self._big_char = [chr(i) for i in range(65, 91)] * mul_num
        self._small_char = [chr(i) for i in range(97, 123)] * mul_num
        if big_char:
            self.choice_list.extend(self._big_char)
        if small_char:
            self.choice_list.extend(self._small_char)

    def get_code_pic(self, out_path: str):
        captcha_image = captcha(
            width=self.width,
            height=self.height,
            drawings=[
                background(),
                text(
                    fonts=[
                        "BRADHITC.TTF",
                    ],
                    drawings=[warp(), rotate(80), offset()],
                    font_sizes=(150, 100, 150)
                ),
                # curve(),
                noise(),
                smooth(),
            ]
        )
        for num, _code in enumerate(self.choice_list):
            image = captcha_image(_code)
            image.save(f"{out_path}/{num}.{_code}.jpg", "JPEG", quality=75)
            print(f"\r已生成：{num + 1}张字符为：{_code}的图片", end='', flush=True)

    @property
    def classes(self) -> int:
        """
        获取类别数

        :return: 类别数
        """
        return len(set(self.choice_list))


if __name__ == '__main__':
    t = Captcha_Product(big_char=True, small_char=True, mul_num=1000)
    t.get_code_pic('../captcha_data')
